import os
from typing import Optional
from langchain.chains import LLMChain
from langchain_core.prompts import PromptTemplate
from src.core.memory import RedisMemory
from src.core.knowledge_base import KnowledgeBase
from src.core.database import DatabaseManager, MySQLConfig

# 从环境变量获取配置
REDIS_URL = os.getenv("REDIS_URL", "redis://localhost:6379")
DB_CONFIG = MySQLConfig(
    host=os.getenv("DB_HOST", "localhost"),
    port=int(os.getenv("DB_PORT", "3306")),
    user=os.getenv("DB_USER", "root"),
    password=os.getenv("DB_PASSWORD", ""),
    database=os.getenv("DB_NAME", "financial_data")
)

# 基本面研究Agent的提示模板
FUNDAMENTAL_PROMPT = PromptTemplate(
    input_variables=["symbol", "financial_data", "industry_context"],
    template="""
    请对{symbol}进行基本面分析：
    
    财务数据：
    {financial_data}
    
    行业背景：
    {industry_context}
    
    请从以下角度进行分析：
    1. 财务健康状况（资产负债表、利润表、现金流量表）
    2. 关键财务比率（PE、PB、ROE、负债率等）
    3. 盈利能力与成长性分析
    4. 估值分析（DCF模型、相对估值法）
    """
)

# 行业研究Agent的提示模板
INDUSTRY_PROMPT = PromptTemplate(
    input_variables=["industry", "market_data", "competitive_analysis"],
    template="""
    请对{industry}行业进行深度研究：
    
    市场数据：
    {market_data}
    
    竞争分析：
    {competitive_analysis}
    
    请从以下角度进行分析：
    1. 行业生命周期阶段
    2. 波特五力竞争模型分析
    3. 产业链结构分析
    4. 行业发展趋势与前景
    """
)

# 宏观研究Agent的提示模板
MACRO_PROMPT = PromptTemplate(
    input_variables=["economic_indicators", "monetary_policy", "global_factors"],
    template="""
    请进行宏观经济分析：
    
    经济指标：
    {economic_indicators}
    
    货币政策：
    {monetary_policy}
    
    国际因素：
    {global_factors}
    
    请从以下角度进行分析：
    1. 主要经济指标分析（GDP、CPI、PPI、PMI等）
    2. 货币政策影响评估
    3. 国际形势对国内市场的影响
    4. 对各行业的影响预测
    """
)

# 市场观察员Agent的提示模板
MARKET_OBSERVER_PROMPT = PromptTemplate(
    input_variables=["market_data", "news_sentiment", "technical_indicators"],
    template="""
    请进行市场观察分析：
    
    市场数据：
    {market_data}
    
    新闻情绪：
    {news_sentiment}
    
    技术指标：
    {technical_indicators}
    
    请从以下角度进行分析：
    1. 实时市场动态监控
    2. 异常交易行为检测
    3. 市场情绪分析
    4. 短期市场趋势预测
    """
)

class FundamentalResearchAgent:
    """负责基本面研究的Agent"""
    def __init__(self, llm):
        self.memory = RedisMemory(redis_url=REDIS_URL)
        self.knowledge_base = KnowledgeBase()
        self.db = DatabaseManager(config=DB_CONFIG)
        self.llm_chain = LLMChain(llm=llm, prompt=FUNDAMENTAL_PROMPT)
    
    def analyze_fundamentals(self, symbol: str) -> dict:
        """分析公司基本面"""
        # 1. 从数据库获取财务数据
        financial_query = f"""
        SELECT * FROM financial_statements 
        WHERE symbol = '{symbol}' 
        ORDER BY report_date DESC 
        LIMIT 4
        """
        financial_data = self.db.execute_query(financial_query)
        
        # 2. 从知识库获取行业背景
        industry_results = self.knowledge_base.search(f"{symbol} 所属行业分析")
        industry_context = "\n".join([res['document'] for res in industry_results])
        
        # 3. 使用LLM进行基本面分析
        analysis = self.llm_chain.run(
            symbol=symbol,
            financial_data=str(financial_data),
            industry_context=industry_context
        )
        
        # 4. 存储分析结果到记忆
        self.memory.store(f"fundamental_analysis:{symbol}", analysis)
        
        return {
            "symbol": symbol,
            "analysis": analysis,
            "timestamp": "current"
        }

class IndustryResearchAgent:
    """负责行业研究的Agent"""
    def __init__(self, llm):
        self.memory = RedisMemory(redis_url=REDIS_URL)
        self.knowledge_base = KnowledgeBase()
        self.db = DatabaseManager(config=DB_CONFIG)
        self.llm_chain = LLMChain(llm=llm, prompt=INDUSTRY_PROMPT)
    
    def analyze_industry(self, industry: str) -> dict:
        """分析行业情况"""
        # 1. 从数据库获取市场数据
        market_query = f"""
        SELECT * FROM industry_data 
        WHERE industry_name = '{industry}' 
        ORDER BY date DESC 
        LIMIT 10
        """
        market_data = self.db.execute_query(market_query)
        
        # 2. 从知识库获取竞争分析
        competitive_results = self.knowledge_base.search(f"{industry} 竞争格局 波特五力")
        competitive_analysis = "\n".join([res['document'] for res in competitive_results])
        
        # 3. 使用LLM进行行业分析
        analysis = self.llm_chain.run(
            industry=industry,
            market_data=str(market_data),
            competitive_analysis=competitive_analysis
        )
        
        # 4. 存储分析结果到记忆
        self.memory.store(f"industry_analysis:{industry}", analysis)
        
        return {
            "industry": industry,
            "analysis": analysis,
            "timestamp": "current"
        }

class MacroResearchAgent:
    """负责宏观研究的Agent"""
    def __init__(self, llm):
        self.memory = RedisMemory(redis_url=REDIS_URL)
        self.knowledge_base = KnowledgeBase()
        self.db = DatabaseManager(config=DB_CONFIG)
        self.llm_chain = LLMChain(llm=llm, prompt=MACRO_PROMPT)
    
    def analyze_macro(self) -> dict:
        """分析宏观经济"""
        # 1. 从数据库获取经济指标
        indicators_query = """
        SELECT * FROM economic_indicators 
        ORDER BY date DESC 
        LIMIT 20
        """
        economic_indicators = self.db.execute_query(indicators_query)
        
        # 2. 从知识库获取货币政策信息
        monetary_results = self.knowledge_base.search("货币政策 央行公告")
        monetary_policy = "\n".join([res['document'] for res in monetary_results])
        
        # 3. 从知识库获取国际因素
        global_results = self.knowledge_base.search("国际贸易 形势分析")
        global_factors = "\n".join([res['document'] for res in global_results])
        
        # 4. 使用LLM进行宏观分析
        analysis = self.llm_chain.run(
            economic_indicators=str(economic_indicators),
            monetary_policy=monetary_policy,
            global_factors=global_factors
        )
        
        # 5. 存储分析结果到记忆
        self.memory.store("macro_analysis", analysis)
        
        return {
            "analysis": analysis,
            "timestamp": "current"
        }

class MarketObserverAgent:
    """负责市场观察的Agent"""
    def __init__(self, llm):
        self.memory = RedisMemory(redis_url=REDIS_URL)
        self.knowledge_base = KnowledgeBase()
        self.db = DatabaseManager(config=DB_CONFIG)
        self.llm_chain = LLMChain(llm=llm, prompt=MARKET_OBSERVER_PROMPT)
    
    def observe_market(self, symbol: Optional[str] = None) -> dict:
        """观察市场动态"""
        # 1. 从数据库获取市场数据
        if symbol:
            market_query = f"""
            SELECT * FROM market_data 
            WHERE symbol = '{symbol}' 
            ORDER BY timestamp DESC 
            LIMIT 50
            """
        else:
            market_query = """
            SELECT * FROM market_data 
            ORDER BY timestamp DESC 
            LIMIT 100
            """
        market_data = self.db.execute_query(market_query)
        
        # 2. 从知识库获取新闻情绪
        news_query = f"{symbol} 新闻情绪分析" if symbol else "市场新闻情绪分析"
        news_results = self.knowledge_base.search(news_query)
        news_sentiment = "\n".join([res['document'] for res in news_results])
        
        # 3. 从知识库获取技术指标
        tech_query = f"{symbol} 技术分析" if symbol else "技术指标分析"
        tech_results = self.knowledge_base.search(tech_query)
        technical_indicators = "\n".join([res['document'] for res in tech_results])
        
        # 4. 使用LLM进行市场观察分析
        analysis = self.llm_chain.run(
            market_data=str(market_data),
            news_sentiment=news_sentiment,
            technical_indicators=technical_indicators
        )
        
        # 5. 存储分析结果到记忆
        key = f"market_observation:{symbol}" if symbol else "market_observation"
        self.memory.store(key, analysis)
        
        return {
            "symbol": symbol,
            "analysis": analysis,
            "timestamp": "current"
        }
# 风控师Agent的提示模板
RISK_PROMPT = PromptTemplate(
    input_variables=["analysis_results", "risk_framework", "market_context"],
    template="""
    请对以下分析结果进行风险评估：
    
    分析结果：
    {analysis_results}
    
    风险评估框架：
    {risk_framework}
    
    市场背景：
    {market_context}
    
    请从以下角度进行风险评估：
    1. 市场风险（波动性、流动性风险）
    2. 信用风险（交易对手风险）
    3. 操作风险（系统风险、人为错误）
    4. 合规风险（监管政策变化）
    5. 风险预警和应对建议
    """
)

# 顾问Agent的提示模板
ADVISOR_PROMPT = PromptTemplate(
    input_variables=["fundamental_analysis", "industry_analysis", "macro_analysis", "risk_assessment"],
    template="""
    请基于以下分析结果生成投资建议：
    
    基本面分析：
    {fundamental_analysis}
    
    行业分析：
    {industry_analysis}
    
    宏观分析：
    {macro_analysis}
    
    风险评估：
    {risk_assessment}
    
    请从以下角度生成投资建议：
    1. 投资评级（买入/持有/卖出）
    2. 目标价格和时间周期
    3. 仓位管理建议
    4. 风险控制措施
    5. 投资策略和执行计划
    """
)

class RiskAgent:
    """负责风险评估的Agent"""
    def __init__(self, llm):
        self.memory = RedisMemory(redis_url=REDIS_URL)
        self.knowledge_base = KnowledgeBase()
        self.db = DatabaseManager(config=DB_CONFIG)
        self.llm_chain = LLMChain(llm=llm, prompt=RISK_PROMPT)
    
    def assess_risk(self, analysis_results: dict) -> dict:
        """评估投资风险"""
        # 1. 从知识库获取风险评估框架
        risk_results = self.knowledge_base.search("投资风险评估框架")
        risk_framework = "\n".join([res['document'] for res in risk_results])
        
        # 2. 从知识库获取市场背景
        market_results = self.knowledge_base.search("当前市场环境分析")
        market_context = "\n".join([res['document'] for res in market_results])
        
        # 3. 使用LLM进行风险评估
        assessment = self.llm_chain.run(
            analysis_results=str(analysis_results),
            risk_framework=risk_framework,
            market_context=market_context
        )
        
        # 4. 存储风险评估结果
        self.memory.store("risk_assessment", assessment)
        
        return {
            "assessment": assessment,
            "timestamp": "current"
        }

class AdvisorAgent:
    """负责投资建议的Agent"""
    def __init__(self, llm):
        self.memory = RedisMemory(redis_url=REDIS_URL)
        self.knowledge_base = KnowledgeBase()
        self.db = DatabaseManager(config=DB_CONFIG)
        self.llm_chain = LLMChain(llm=llm, prompt=ADVISOR_PROMPT)
    
    def provide_advice(self, analyses: dict) -> dict:
        """生成投资建议"""
        # 1. 从各个分析结果中提取信息
        fundamental_analysis = analyses.get("fundamental", {}).get("analysis", "")
        industry_analysis = analyses.get("industry", {}).get("analysis", "")
        macro_analysis = analyses.get("macro", {}).get("analysis", "")
        risk_assessment = analyses.get("risk", {}).get("assessment", "")
        
        # 2. 使用LLM生成投资建议
        advice = self.llm_chain.run(
            fundamental_analysis=fundamental_analysis,
            industry_analysis=industry_analysis,
            macro_analysis=macro_analysis,
            risk_assessment=risk_assessment
        )
        
        # 3. 存储投资建议
        self.memory.store("investment_advice", advice)
        
        return {
            "advice": advice,
            "timestamp": "current"
        }